import concurrent.futures
import logging
import time
from datetime import datetime, timedelta
from typing import List, Optional
import akshare as ak
import pandas as pd
from sqlalchemy.ext.declarative import declarative_base
from tqdm import tqdm

from infrastructure.model.stock_kline import StockKline
from utils.logger import ConsoleLogger

Base = declarative_base()

# 创建控制台日志记录器
logger = ConsoleLogger(
    name="akshare_provider",
    level=logging.INFO
)


class AkshareProvider:
    def __init__(self, max_workers: int = 10,
                 retry_count: int = 3, delay: float = 0.1):
        """
        Initialize MarketDataFetcher with PostgreSQL connection

        Args:
            max_workers: Maximum number of worker threads
            retry_count: Number of retries for failed requests
            delay: Delay between retries in seconds
        """
        self.max_workers = max_workers
        self.retry_count = retry_count
        self.delay = delay

    def _clean_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """Clean and validate market data"""
        if df.empty:
            return df

        # Convert date column, yyyy-MM-dd

        df['trade_date'] = pd.to_datetime(df['日期']).dt.strftime('%Y-%m-%d')
        # Rename columns to match database schema
        column_mapping = {
            '开盘': 'open',
            '收盘': 'close',
            '最高': 'high',
            '最低': 'low',
            '成交量': 'volume'
        }
        df = df.rename(columns=column_mapping)

        # Select and reorder columns
        columns = ['trade_date', 'open', 'high', 'low', 'close', 'volume']
        df = df[columns]

        # Remove rows with missing values
        df = df.dropna()

        # Remove extreme values (e.g., prices <= 0)
        df = df[df['close'] > 0]

        # Convert numeric columns to appropriate types
        for col in ['open', 'high', 'low', 'close']:
            df[col] = pd.to_numeric(df[col], errors='coerce').round(2)
        df['volume'] = pd.to_numeric(df['volume'], errors='coerce').astype('int64')

        return df

    @classmethod
    def get_stock_codes(cls) -> List[str]:
        """Get list of all A-stock codes"""
        try:
            stock_info_df = ak.stock_info_a_code_name()
            return [code for code in stock_info_df['code'].tolist()]
        except Exception as e:
            logger.info(f"Failed to get stock codes: {str(e)}")
            return []

    def _fetch_stock_data(self, ts_code: str, end_date: datetime,
                          start_date: Optional[datetime] = None) -> Optional[pd.DataFrame]:
        """Fetch stock data for a given symbol and date range, fetching backwards from end_date in 500-day segments"""
        latest_date = StockKline.get_latest_date(ts_code)

        symbol = ts_code.split('.')[0]  # Remove exchange suffix for akshare

        if latest_date is not None:
            if end_date <= latest_date:
                logger.info(f"{ts_code}: Data is up to date")
                return None
            start_date = latest_date + timedelta(days=1)
        elif start_date is None:
            start_date = datetime(1980, 12, 19)

        logger.info(f"{ts_code}: Fetching data from {start_date.date()} to {end_date.date()} starting")

        all_data = []
        current_end = end_date
        segment_size = timedelta(days=1000)

        while current_end >= start_date:
            current_start = max(current_end - segment_size, start_date)
            for attempt in range(self.retry_count):
                try:
                    df = ak.stock_zh_a_hist(symbol=symbol,
                                            start_date=current_start.strftime('%Y%m%d'),
                                            end_date=current_end.strftime('%Y%m%d'),
                                            adjust='qfq')
                    if df is not None and not df.empty:
                        df = self._clean_data(df)
                        if df is not None and not df.empty:
                            df['ts_code'] = ts_code
                            all_data.append(df)
                            break
                    time.sleep(self.delay)
                except Exception as e:
                    logger.info(
                        f"Attempt {attempt + 1} failed for {ts_code} from {current_start.date()} to {current_end.date()}: {str(e)}")
                    if attempt < self.retry_count - 1:  # 如果不是最后一次重试，则等待后重试
                        time.sleep(self.delay * (attempt + 1))
                    else:  # 如果是最后一次重试失败，打印警告并继续下一个时间段
                        logger.warning(
                            f"Warning: Failed to fetch data for period {current_start.date()} to {current_end.date()}")

            current_end = current_start - timedelta(days=1)

        logger.info(f"{ts_code}: Fetching data from {start_date.date()} to {end_date.date()} end")
        if all_data:
            # 合并数据并按日期排序
            result = pd.concat(all_data, ignore_index=True)
            return result.sort_values('trade_date').reset_index(drop=True)
        return None

    def update_stock_data(self, symbols: Optional[List[str]] = None,
                          end_date: Optional[datetime] = None) -> None:
        """Update stock data for given symbols"""
        if symbols is None:
            symbols = self.get_stock_codes()

        if end_date is None:
            end_date = datetime.now()

        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = []
            for symbol in symbols:
                future = executor.submit(self._fetch_stock_data, symbol, end_date)
                futures.append((symbol, future))

            for symbol, future in tqdm(futures, desc="Updating stock data"):
                try:
                    df = future.result()
                    if df is not None:
                        StockKline.save(df, index_elements=['ts_code', 'trade_date'])
                except Exception as e:
                    logger.info(f"Failed to process {symbol}: {str(e)}")

    def sync_realtime_market_data(self):
        """
        同步所有市场的实时行情数据到 StockKline 表
        包括：深圳、上海、北京、创业板、科创板等市场，
        每天行情结束后
        """
        result = self.get_realtime_market_data()

        # 只保留需要的列
        result = result[['ts_code', 'open', 'high', 'low', 'close', 'volume']]

        StockKline.save(result, index_elements=['ts_code', 'trade_date'])

    def get_realtime_market_data(self) -> pd.DataFrame:
        """
            同步所有市场的实时行情数据到 StockKline 表
            包括：深圳、上海、北京、创业板、科创板等市场，
            每天行情结束后
        """
        all_market_data = []

        # 定义所有需要获取的市场及其对应的函数
        market_functions = {
            'SZ': ak.stock_sz_a_spot_em,  # 深圳A股
            'BJ': ak.stock_bj_a_spot_em,  # 北京A股
            'CY': ak.stock_cy_a_spot_em,  # 创业板
            'KC': ak.stock_kc_a_spot_em,  # 科创板
            'NEW': ak.stock_new_a_spot_em  # 新股
        }

        current_time = pd.Timestamp.now().strftime('%Y-%m-%d')

        for market, func in market_functions.items():
            try:
                logger.info(f"Fetching {market} market data...")
                df = func()
                if df is not None and not df.empty:
                    # 统一字段名称
                    column_mapping = {
                        '代码': 'ts_code',  # 股票代码
                        '名称': 'name',  # 股票名称
                        '最新价': 'close',  # 收盘价
                        '涨跌幅': 'pct_change',  # 涨跌幅（百分比变化）
                        '涨跌额': 'change',  # 涨跌额
                        '成交量': 'volume',  # 成交量
                        '成交额': 'turnover',  # 成交金额
                        '振幅': 'amplitude',  # 振幅
                        '最高': 'high',  # 最高价
                        '最低': 'low',  # 最低价
                        '今开': 'open',  # 今日开盘价
                        '昨收': 'prev_close',  # 昨日收盘价
                        '量比': 'volume_ratio',  # 量比
                        '换手率': 'turnover_rate',  # 换手率
                        '市盈率-动态': 'pe_dynamic',  # 动态市盈率
                        '市净率': 'pb_ratio',  # 市净率
                        '总市值': 'total_market_cap',  # 总市值
                        '流通市值': 'circulating_market_cap',  # 流通市值
                        '涨速': 'price_change_rate',  # 涨速
                        '5分钟涨跌': 'five_min_change',  # 5分钟涨跌
                        '60日涨跌幅': 'sixty_day_change',  # 60日涨跌幅
                        '年初至今涨跌幅': 'ytd_change'  # 年初至今涨跌幅
                    }

                    # 重命名列
                    df = df.rename(columns=lambda x: column_mapping.get(x, x))

                    df.dropna(how='any', inplace=True)

                    # 将数值类型的列转换为数值
                    numeric_columns = ['open', 'high', 'low', 'close']
                    for col in numeric_columns:
                        if col in df.columns:
                            df[col] = pd.to_numeric(df[col], errors='coerce').round(2)

                    df['volume'] = pd.to_numeric(df['volume'], errors='coerce').astype('int64')

                    # 添加交易日期
                    df['trade_date'] = current_time

                    all_market_data.append(df)
                    logger.info(f"Successfully fetched {market} market data: {len(df)} records")

                time.sleep(self.delay)  # 避免请求过快

            except Exception as e:
                logger.info(f"Error fetching {market} market data: {str(e)}")
                continue

        if not all_market_data:
            return pd.DataFrame()

        # 合并所有市场数据
        result = pd.concat(all_market_data, ignore_index=True)
        return result
